# 使用 RNN 处理语句，引入embedding

import torch

# parameters
num_class = 4
input_size = 4
hidden_size = 8
embedding_size = 10
num_layers = 2
# batch_size 只用于生成初始的隐层输入 h0 有关
batch_size = 1
seq_len = 5

# 词汇表
idx2char = ['e', 'h', 'l', 'o']
# (batch, seq_len)
x_data = [[1, 0, 2, 2, 3]]
# (batch * seq_len)
y_data = [3, 1, 2, 3, 2]

inputs = torch.LongTensor(x_data)
labels = torch.LongTensor(y_data)


class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        # 定义Embedding操作，输入的样本维度，embedding维度
        self.emb = torch.nn.Embedding(input_size, embedding_size)
        # 定义 RNN，输入维度（batchSize,seqLen, embeddingSize），输出维度（batchSize,seqLen, hiddenSize）
        self.rnn = torch.nn.RNN(input_size=embedding_size,
                                hidden_size=hidden_size,
                                num_layers=num_layers,
                                # batch_first=True 表明定义维度时batchSize放在前面，如果不说明则是（seqLen,batchSize,embeddingSize）
                                batch_first=True)
        # 定义一个全连接层，in_features 为 hidden_size 隐层特征维度，out_features 为 num_class 标签的维度
        self.fc = torch.nn.Linear(hidden_size, num_class)

    def forward(self, x):
        # 初始化隐层输入h0
        hidden = torch.zeros(num_layers, x.size(0), hidden_size)
        # (batch, seqLen, embeddingSize)
        x = self.emb(x)
        # 获得最后的输出和隐层输出hidden
        x, _ = self.rnn(x, hidden)
        # 将RNN得到的输出接到全连接层
        x = self.fc(x)
        # 输出分类结果
        return x.view(-1, num_class)


# 模型实例
model = Model()

# 定义损失和优化器
criterion = torch.nn.CrossEntropyLoss()
# 这里改用一个新的算法 Adam 作为优化算法，学习率改为0.1加速收敛速度
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

# 训练模型
for epoch in range(15):
    # 梯度清0
    optimizer.zero_grad()
    # 获得输出
    outputs = model(inputs)
    # 计算损失
    loss = criterion(outputs, labels)
    # 反向传播
    loss.backward()
    # 参数更新
    optimizer.step()
    _, idx = outputs.max(dim=1)
    idx = idx.data.numpy()
    print('Predicted: ', ''.join([idx2char[x] for x in idx]), end='')
    print(', Epoch [%d/15] loss = %.3f' % (epoch + 1, loss.item()))
